package gshutdown

import (
	"context"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"gitee.com/youkelike/zlog"
)

// 回调定义
type ShutdownCallback interface {
	OnShutdown() error
}

type ShutdownFunc func() error

func (f ShutdownFunc) OnShutdown() error {
	return f()
}

// 错误定义
type ErrHandler interface {
	OnError(err error)
}

type ErrFunc func(err error)

func (f ErrFunc) OnError(err error) {
	f(err)
}

// 优雅关停功能定义
type GSInterface interface {
	SetErrHandler(e ErrHandler)
	ReportError(err error)
	AddShutdownCallback(f ShutdownCallback)

	Start()
	Wait(ctx context.Context)
}

type GracefulShutdown struct {
	callbacks    []ShutdownCallback
	errorHandler ErrHandler
	wg           sync.WaitGroup
}

func New() *GracefulShutdown {
	return &GracefulShutdown{
		callbacks: make([]ShutdownCallback, 0, 10),
	}
}

func (gs *GracefulShutdown) Start() {
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
		<-c

		gs.StartShutdown()
	}()
}

func (gs *GracefulShutdown) AddShutdownCallback(fn ShutdownCallback) {
	gs.callbacks = append(gs.callbacks, fn)
}

func (gs *GracefulShutdown) SetErrHandler(e ErrHandler) {
	gs.errorHandler = e
}

func (gs *GracefulShutdown) StartShutdown() {
	for _, cb := range gs.callbacks {
		go func(fn ShutdownCallback) {
			defer gs.wg.Done()

			gs.ReportError(fn.OnShutdown())
		}(cb)
	}
}
func (gs *GracefulShutdown) Wait(ctx context.Context) {
	// 根据 callbacks 数量提前设置 wg
	gs.wg.Add(len(gs.callbacks))
	zlog.Infow("callback length: %d", len(gs.callbacks))

	ch := make(chan struct{}, 1)
	go func() {
		gs.wg.Wait()
		ch <- struct{}{}
	}()
	select {
	case <-ch:
	case <-ctx.Done():
		gs.ReportError(ctx.Err())
	}

	zlog.Infow("graceful shutdown finish")
}

func (gs *GracefulShutdown) ReportError(err error) {
	if err == nil || gs.errorHandler == nil {
		return
	}
	gs.errorHandler.OnError(err)
}
